Pytorch 多进程在单卡上测试 您所在的位置:网站首页 python 多线程 GPU Pytorch 多进程在单卡上测试

Pytorch 多进程在单卡上测试

2024-01-31 18:04| 来源: 网络整理| 查看: 265

有些炼丹师可能机器不足,只有一张卡,然后训练完成了,想要测试的时候,受限于图像样本size不一致,不能合并到一个batch中。当然解决方案很多,但有一种更自然的办法,既使用任意分辨率的输入,同样使用多进程在单卡上并行执行数据的测试,从而加快测试速度,赶上dideline。

大家都知道Pytorch官方推广大家使用分布式多卡并行计算,其原理是每个进程都使用一个GPU。那我们现在要做的就是多进程都在一个GPU上,每个进程处理一批数据,从而加快处理速度。

前言

模型在GPU上,Pytorch规定:多进程的启动方法必须是"spawn",所以使用map_async或者apply_async这类方法是不行了。

执行这套流程的大致思路是,我们可以先把要处理的数据加入到容器中(比如字典或者元祖列表),然后把容器中的数据分段,分段的数目就是进程数,然后每个进程处理自己对应的那些数据。由于输入的size不一样,所以我们要一张一张的处理,但因为开了多进程,实际的batchsize等于进程数目。

假设我们现在有一批数据,我们先要把这批数据的目录地址加到一个容器中,

target_loc = make_dataset(os.path.join(img_root, training_subfix), True) target_loc.extend(make_dataset(os.path.join(img_root, val_subfix), False)) target_loc.sort(key=lambda x: x[0]) # 保证一致 print(len(target_loc))

make_dataset的作用就是把所有的数据的地址加入到target_loc中。target_loc是一个list。 然后开始多进程启动

import multiprocessing as mp # 这里不需要使用pytorch的multiprocessing mp = mp.get_context('spawn') pool = [] n_proc = 4 # 开4进程 total_num = len(target_loc) phrase = total_num // n_proc + 1 for i in range(n_proc): divied_target = target_loc[i*phrase:(i+1)*phrase] # 所有的数据分成4份 process = mp.Process(target=prediction_by_dfsd, args=(divied_target,i)) # 每个进程都执行prediction_by_dfsd process.start() pool.append(process) for p in pool: p.join() # 等待所有进程执行完毕

prediction_by_dfsd这个函数就是测试数据的函数,这个函数需要完成的是,从分段的容器中读取图像,然后一张张的预测,然后保存结果。 另外我需要说明的是,spawn的启动进程的方式和我们平常使用的map_async等不一样,从spawn启动的进程,是无法访问到主进程的任何变量的。所有我们需要在prediction_by_dfsd完成模型定义,导入参数等操作。当然一些变量可以通过Process的args传进去。 所以和使用分布式训练一样,每个进程都会定义一次模型。只不过我们这个方法,是在单卡上使用多进程,一样能使用并行预测,满足了在一些情况,比如样本size不一样,不能合并到batch中,而一张一张预测又太慢的问题。

最后

当然我后面反应过来了。只用使用dataloader,batchsize为1,一样可以使用标准的pytorch分布式流程使用多进程在单卡上的测试,只不过不需要设置gpu id就行了。所有的模型都往一个gpu上搬运。不过呢,对于简单的任务,没有必要再写一个dataloader,所以我这个办法还是有价值的。



【本文地址】

公司简介

联系我们

今日新闻

    推荐新闻

    专题文章
      CopyRight 2018-2019 实验室设备网 版权所有